import torch
import torch.nn as nn
import torch.nn.functional as F
# from utils.soft_skeleton import soft_skel


class dice_bce_loss(nn.Module):
    """
    道路分割的损失： 交叉熵损失 + dice loss
    总损失： 初步分割的损失 + 精细分割的损失

    """
    def __init__(self, batch=True):
        super(dice_bce_loss, self).__init__()
        self.batch = batch
        self.bce_loss = nn.BCELoss()

    def soft_dice_coeff(self, y_true, y_pred):
        smooth = 0.0  # may change
        if self.batch:
            i = torch.sum(y_true)
            j = torch.sum(y_pred)
            intersection = torch.sum(y_true * y_pred)
        else:
            i = y_true.sum(1).sum(1).sum(1)
            j = y_pred.sum(1).sum(1).sum(1)
            intersection = (y_true * y_pred).sum(1).sum(1).sum(1)
        score = (2. * intersection + smooth) / (i + j + smooth)
        # score = (intersection + smooth) / (i + j - intersection + smooth)#iou
        return score.mean()

    def soft_dice_loss(self, y_true, y_pred):
        loss = 1 - self.soft_dice_coeff(y_true, y_pred)
        return loss

    def __call__(self, y_pred, y_true):
        a = self.bce_loss(y_pred, y_true)
        b = self.soft_dice_loss(y_true, y_pred)
        return a + b
    

def dice_bce_loss_func(y_pred, y_true):
    """
    Segmentation loss
    dice loss + bce loss
    """
    a = F.binary_cross_entropy(y_pred, y_true)

    smooth = 0.0  # may change
    i = torch.sum(y_true)
    j = torch.sum(y_pred)
    intersection = torch.sum(y_true * y_pred)
    score = (2. * intersection + smooth) / (i + j + smooth)
    b = 1 - score.mean()

    return a + b


def con_loss_func(y_pred, y_true):
    """
    Connectivity loss
    bce loss
    """
    loss = F.binary_cross_entropy_with_logits(y_pred, y_true)
    return loss
    

# def cldice_bce_loss_func(y_pred, y_true, dataset='DeepGlobe', smooth=1e-5):
#     """
#     cldice loss(A topology-preserving loss) + bce loss
#     """
#     if dataset == 'DeepGlobe':  # DeepGlobe
#         iter_num = 51
#     elif dataset == 'LowContrast':  # Low Contrast Roads
#         iter_num = 35
#     else:  # Massachusetts
#         iter_num = 25
#
#     bce = F.binary_cross_entropy(y_pred, y_true)
#
#     skel_pred = soft_skel(y_pred, iter_num)
#     skel_true = soft_skel(y_true, iter_num)
#     tprec = (torch.sum(torch.multiply(skel_pred, y_true)) + smooth) / (
#                 torch.sum(skel_pred) + smooth)
#     tsens = (torch.sum(torch.multiply(skel_true, y_pred)) + smooth) / (
#                 torch.sum(skel_true) + smooth)
#     cl_dice = 1. - 2.0 * (tprec * tsens) / (tprec + tsens)
#
#     return bce + cl_dice


